#include	"task.h"
#include	"paging.h"

#include	"monitor.h"

extern struct page_directory *kernel_directory;
extern struct page_directory *current_directory;
extern u32int initial_esp;
extern u32int read_eip();

extern u8int sched_reentry;

u32int next_pid = 0;

volatile struct task *current_task = 0;
volatile struct task *ready_queue = 0;
volatile struct task *sched_task = 0;
volatile struct task *current_task_bak = 0;


void init_tasking(){

	asm volatile("cli");

	/* assume a stack at 0xE0000000 of length 0x2000) */
	move_stack((void *)0xE0000000, 0x2000);

	/* craft kernel task */
	current_task = ready_queue = (struct task *)kmalloc(sizeof(struct task));
	current_task -> id = next_pid++;
	current_task -> esp = current_task -> ebp = 0;
	current_task -> eip = 0;
	current_task -> page_dir = current_directory;
	current_task -> quantum = DEFAULT_QUANTA;
	current_task -> status = TASK_READY;
	current_task -> next = 0;
	current_task -> kernel_stack = kmalloc_a(KERNEL_STACK_SIZE);

	monitor_write("kernel stack @ ");
	monitor_write_hex(current_task -> kernel_stack);
	monitor_write("\n");

	asm volatile("sti");
}

static void idle(){
	current_task -> status = TASK_BLOCK;
}

void sched(){
	monitor_write("[sched] Scheduling @ task ");
	monitor_write_hex(current_task -> id);
	monitor_write("\n");

	task_switch(1);
	while(1){
		//task_switch(1);
		idle();
	}
}

void task_switch(int flag_force){

	/* when switching task, we always using ppda of task 0 */

	u32int eip, ebp, esp;

	/* handles timer irq before we initialized multitasking */
	if(!current_task)
		return;

	sched_reentry = 3;

	asm volatile("mov %%esp, %0" : "=r"(esp));
	asm volatile("mov %%ebp, %0" : "=r"(ebp));

	eip = read_eip();

	current_task -> eip = eip;
	current_task -> ebp = ebp;
	current_task -> esp = esp;
/*
	monitor_write("[swtch] eip = ");
	monitor_write_hex(eip);
	monitor_write("\n");
*/	
	while(1){
		current_task = current_task -> next;
		if(!current_task)
			break;
		if(current_task -> status == TASK_READY)
			break;
	}
	if(!current_task){
		current_task = ready_queue;
		current_task -> status = TASK_READY;
	}
	esp = current_task -> esp;
	ebp = current_task -> ebp;
	eip = current_task -> eip;

	current_directory = current_task -> page_dir;
	set_kernel_stack(current_task -> kernel_stack + KERNEL_STACK_SIZE);

	monitor_write("[swtch] kernel stack @ ");
	monitor_write_hex(current_task -> kernel_stack + KERNEL_STACK_SIZE);
	monitor_write("\n");

	if(current_task -> id > 0){
		monitor_write("[swtch] child stub - ctask.eip = ");
		monitor_write_hex(eip);
		monitor_write(", ctask.id = ");
		monitor_write_hex(current_task -> id);
		monitor_write(", status = ");
		monitor_write((current_task -> status == TASK_READY) ? "ready" : "block");
		monitor_write("\n");
		//while(1);
	}
/*
	if(!flag_force){
		monitor_write("IRQ switching.\n");
		kfree(current_task);
	}
	current_task = sched_task;
*/	
	asm volatile("		\
			cli;	\
			mov %0, %%ecx;	\
			mov %1, %%esp;	\
			mov %2, %%ebp;	\
			mov %3, %%cr3;	\
			mov %%ecx, %%eax; \
			sti;	\
			jmp *%%ecx"
			: : "r"(eip), "r"(esp), "r"(ebp), "r"(current_directory -> physaddr));
}


void move_stack(void *new_stack_start, u32int size){

	u32int i;
	u32int pd_addr, old_stack_pointer, old_base_pointer,
		new_stack_pointer, new_base_pointer, offset, tmp, *tmp2;

	for(i = (u32int)new_stack_start; i >= (u32int)new_stack_start - size; i -= 0x1000){
		alloc_frame(get_page(i, 1, current_directory), 0, 1);
	}
	asm volatile("mov %%cr3, %0" : "=r" (pd_addr));
	asm volatile("mov %0, %%cr3" : : "r" (pd_addr));

	asm volatile("mov %%esp, %0" : "=r" (old_stack_pointer));
	asm volatile("mov %%ebp, %0" : "=r" (old_base_pointer));

	offset = (u32int)new_stack_start - initial_esp;
	new_stack_pointer = old_stack_pointer + offset;
	new_base_pointer = old_base_pointer + offset;
	memcpy((void *)new_stack_pointer, (void *)old_stack_pointer, initial_esp - old_stack_pointer);
	
	/* FATAL BUG - we may override data which is not ebp */
	
	for(i = (u32int)new_stack_start; i > (u32int)new_stack_start - size; i -= 4){
		tmp = *(u32int *)i;
		if((tmp > old_stack_pointer) && (tmp < initial_esp)){
			tmp += offset;
			tmp2 = (u32int *)i;
			*tmp2 = tmp;
		}
	}
	

	asm volatile("mov %0, %%esp" : : "r"(new_stack_pointer));
	asm volatile("mov %0, %%ebp" : : "r"(new_base_pointer));
}


void switch_to_usermode(){

	set_kernel_stack(current_task -> kernel_stack + KERNEL_STACK_SIZE);

	asm volatile(" \
		cli;	\
		mov $0x23, %ax;	\
		mov %ax, %ds;		\
		mov %ax, %es;		\
		mov %ax, %fs;		\
		mov %ax, %gs;		\
		\
		mov %esp, %eax;		\
		pushl $0x23;		\
		pushl %esp;		\
		pushf;			\
		pop %eax;		\
		orl $0x200, %eax;	\
		pushl %eax;		\
		pushl $0x1B;		\
		push $1f;		\
		iret;			\
	1: \
		");
}


int fork(){

	struct task *parent_task, *new_task, *tmp_task;
	struct page_directory *directory;

	u32int eip, esp, ebp;

	/* atomic fork() */
	asm volatile("cli");

	parent_task = (struct task *)current_task;
	directory = clone_directory(current_directory);

	new_task = (struct task *)kmalloc(sizeof(struct task));
	new_task -> id = next_pid++;
	new_task -> esp = new_task -> ebp = 0;
	new_task -> eip = 0;
	new_task -> page_dir = directory;
	new_task -> kernel_stack = kmalloc_a(KERNEL_STACK_SIZE); /* BUG: kstack is not a kernel object */
	new_task -> quantum = DEFAULT_QUANTA;
	new_task -> status = TASK_READY;
	new_task -> next = 0;

	tmp_task = (struct task *)ready_queue;
	while(tmp_task -> next)
		tmp_task = tmp_task -> next;
	tmp_task -> next = new_task;
	
	/* entry point for new task */
	eip = read_eip();
	monitor_write("[fork] eip read = ");
	monitor_write_hex(eip);
	monitor_write((current_task == parent_task) ? " - parent" : " - child");
	monitor_write("\n");

	struct page *tmppage = get_page(eip, 0, current_directory);
	monitor_write("\t @ physical page ");
	monitor_write_hex((u32int)((tmppage -> frame) << 12));
	monitor_write("\n");

	/* fork here */
	if(current_task == parent_task){
		sched_reentry = 4;

		asm volatile("mov %%esp, %0" : "=r"(esp));
		asm volatile("mov %%ebp, %0" : "=r"(ebp));
		new_task -> esp = esp;
		new_task -> ebp = ebp;
		new_task -> eip = eip;
		asm volatile("sti");
		/* here we could be scheduled again */
		return new_task -> id;	
	}else{

		sched_reentry = 2;

		monitor_write("[fork] child, proc #");
		monitor_write_hex(current_task -> id);
		monitor_write("\n");
		/* BUG - A WEIRD RETURN ADDR */
		return 0;
		monitor_write("[fork] fatal error - no return\n");
	}
	/* test stub */
	monitor_write("[fork] fatal error - illegal return\n");
}


int getpid(){
	return current_task -> id;
}
